#!/bin/bash
# ==============================================================================
# GSICoreAnalysis.jl - Visualization Runner Script
# ==============================================================================
# This script runs GSI analysis with comprehensive visualization capabilities,
# generating high-quality plots and analysis reports.

set -e

# Color codes
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
MAGENTA='\033[0;35m'
NC='\033[0m'

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LOG_DIR="$SCRIPT_DIR/logs"
VIS_DIR="$SCRIPT_DIR/drp4dvar_visualizations"
JULIA_THREADS="${JULIA_NUM_THREADS:-4}"

mkdir -p "$LOG_DIR" "$VIS_DIR"

print_banner() {
    echo -e "${MAGENTA}================================================${NC}"
    echo -e "${MAGENTA}  GSICoreAnalysis.jl - Visualization Suite${NC}"
    echo -e "${MAGENTA}================================================${NC}"
    echo "🎨 Atmospheric Data Assimilation Visualization"
    echo ""
    echo "📁 Script Directory: $SCRIPT_DIR"
    echo "📊 Output Directory: $VIS_DIR"
    echo "📝 Log Directory: $LOG_DIR"
    echo "🧵 Julia Threads: $JULIA_THREADS"
    echo ""
}

print_usage() {
    echo "Usage: $0 [OPTIONS] [VIS_TYPE]"
    echo ""
    echo "VISUALIZATION TYPES:"
    echo "  full           - Complete visualization suite (6 plots + report)"
    echo "  quick          - Quick visualization (3 key plots)"
    echo "  convergence    - Algorithm convergence comparison only"
    echo "  spatial        - Spatial observation distribution"
    echo "  stats          - Innovation statistics and quality metrics"
    echo "  performance    - Algorithm performance comparison"
    echo "  custom         - Custom visualization configuration"
    echo ""
    echo "OPTIONS:"
    echo "  -h, --help     - Show this help message"
    echo "  -q, --quality Q - Plot quality: low|medium|high|ultra (default: high)"
    echo "  -f, --format F  - Output format: png|pdf|svg (default: png)"
    echo "  -d, --dpi N     - Plot resolution DPI (default: 300)"
    echo "  -t, --threads N - Julia threads (default: 4)"
    echo "  -o, --output DIR - Custom output directory"
    echo "  -i, --interactive - Enable interactive plots"
    echo "  --no-report    - Skip generating analysis report"
    echo "  --clean        - Clean previous results before running"
    echo "  --backend BACKEND - Plot backend: plotlyjs|gr|pyplot (default: plotlyjs)"
    echo ""
    echo "EXAMPLES:"
    echo "  $0                          # Full visualization suite"
    echo "  $0 quick                    # Quick 3-plot visualization"
    echo "  $0 -q ultra -f pdf full     # Ultra quality PDF output"
    echo "  $0 --interactive spatial    # Interactive spatial plots"
    echo "  $0 --clean performance      # Clean run of performance plots"
    echo ""
}

check_visualization_deps() {
    echo "🔍 Checking visualization dependencies..."

    local deps_to_install=()

    # Check for required packages
    julia --project=. -e "
    using Pkg
    required = [\"Plots\", \"StatsPlots\", \"PlotlyJS\", \"StatsBase\", \"Distributions\"]
    installed = keys(Pkg.project().dependencies)
    missing = [pkg for pkg in required if !(pkg in String.(installed))]
    if !isempty(missing)
        println(\"MISSING: \", join(missing, \", \"))
        exit(1)
    else
        println(\"✓ All visualization dependencies available\")
    end
    " 2>/dev/null || {
        echo "📦 Installing missing visualization packages..."
        julia --project=. -e "
        using Pkg
        Pkg.add([\"Plots\", \"StatsPlots\", \"PlotlyJS\", \"StatsBase\", \"Distributions\"])
        Pkg.precompile()
        " || {
            echo -e "${RED}ERROR: Failed to install visualization dependencies${NC}" >&2
            exit 1
        }
    }

    echo -e "${GREEN}✓ Visualization dependencies ready${NC}"
}

setup_julia_environment() {
    # Set Julia visualization environment
    export JULIA_NUM_THREADS="$JULIA_THREADS"
    export GSI_PLOT_QUALITY="$PLOT_QUALITY"
    export GSI_PLOT_FORMAT="$PLOT_FORMAT"
    export GSI_PLOT_DPI="$PLOT_DPI"
    export GSI_PLOT_BACKEND="$PLOT_BACKEND"
    export GSI_OUTPUT_DIR="$VIS_DIR"
    export GSI_INTERACTIVE="$INTERACTIVE"
    export GSI_GENERATE_REPORT="$GENERATE_REPORT"
}

clean_previous_results() {
    if [[ "$CLEAN_RUN" == "true" ]]; then
        echo "🧹 Cleaning previous visualization results..."
        if [[ -d "$VIS_DIR" ]]; then
            find "$VIS_DIR" -name "gsi_*.png" -o -name "gsi_*.pdf" -o -name "gsi_*.svg" -delete 2>/dev/null || true
            find "$VIS_DIR" -name "GSI_Analysis_Report.md" -delete 2>/dev/null || true
        fi
        echo -e "${GREEN}✓ Previous results cleaned${NC}"
    fi
}

run_visualization() {
    local vis_type="$1"
    local timestamp=$(date +%Y%m%d_%H%M%S)
    local log_file="$LOG_DIR/gsi_visualization_${vis_type}_${timestamp}.log"

    echo -e "${BLUE}🎨 Starting $vis_type visualization...${NC}"
    echo "📝 Log file: $log_file"

    # Determine script and configuration
    local julia_script=""
    local plot_config=""

    case "$vis_type" in
        "full"|"complete")
            julia_script="demos/gsi_test_case_demo.jl"
            plot_config="FULL_SUITE"
            ;;
        "quick")
            julia_script="demos/gsi_test_case_demo.jl"
            plot_config="QUICK_VIS"
            ;;
        "convergence")
            julia_script="demos/performance_demo.jl"
            plot_config="CONVERGENCE_ONLY"
            ;;
        "spatial")
            julia_script="demos/gsi_test_case_demo.jl"
            plot_config="SPATIAL_ONLY"
            ;;
        "stats")
            julia_script="demos/gsi_test_case_demo.jl"
            plot_config="STATS_ONLY"
            ;;
        "performance")
            julia_script="demos/performance_demo.jl"
            plot_config="PERFORMANCE_ONLY"
            ;;
        *)
            echo -e "${RED}ERROR: Unknown visualization type: $vis_type${NC}" >&2
            exit 1
            ;;
    esac

    if [[ ! -f "$julia_script" ]]; then
        echo -e "${RED}ERROR: Script not found: $julia_script${NC}" >&2
        exit 1
    fi

    # Setup environment
    setup_julia_environment
    export GSI_VIS_CONFIG="$plot_config"

    # Execute visualization
    local start_time=$(date +%s)
    echo -e "${GREEN}🚀 Executing: julia --project=. --threads=$JULIA_THREADS $julia_script${NC}"

    julia --project=. --threads="$JULIA_THREADS" "$julia_script" > "$log_file" 2>&1
    local exit_code=$?

    local end_time=$(date +%s)
    local duration=$((end_time - start_time))

    if [[ $exit_code -eq 0 ]]; then
        echo -e "${GREEN}✅ Visualization completed successfully${NC}"
        echo "⏱️  Duration: ${duration}s"

        # Count generated files
        local plot_count=0
        local report_count=0

        if [[ -d "$VIS_DIR" ]]; then
            plot_count=$(find "$VIS_DIR" -name "gsi_*.$PLOT_FORMAT" 2>/dev/null | wc -l)
            report_count=$(find "$VIS_DIR" -name "GSI_Analysis_Report.md" 2>/dev/null | wc -l)
        fi

        echo "📊 Generated files:"
        echo "   • $plot_count visualization plots"
        echo "   • $report_count analysis reports"
        echo "📁 Output directory: $VIS_DIR"

        # List generated files
        if [[ $plot_count -gt 0 ]]; then
            echo ""
            echo "📈 Generated plots:"
            find "$VIS_DIR" -name "gsi_*.$PLOT_FORMAT" 2>/dev/null | sort | sed 's/.*\//   • /'
        fi

        if [[ $report_count -gt 0 ]]; then
            echo ""
            echo "📄 Generated reports:"
            find "$VIS_DIR" -name "GSI_Analysis_Report.md" 2>/dev/null | sed 's/.*\//   • /'
        fi

    else
        echo -e "${RED}❌ Visualization failed with exit code: $exit_code${NC}" >&2
        echo "📝 Check log file: $log_file" >&2
        echo ""
        echo "Last 10 lines of log:"
        tail -n 10 "$log_file" | sed 's/^/   /' >&2
        exit $exit_code
    fi
}

# ==============================================================================
# MAIN SCRIPT
# ==============================================================================

# Default values
VIS_TYPE="full"
PLOT_QUALITY="high"
PLOT_FORMAT="png"
PLOT_DPI="300"
PLOT_BACKEND="plotlyjs"
INTERACTIVE="false"
GENERATE_REPORT="true"
CLEAN_RUN="false"
CUSTOM_OUTPUT=""

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            print_usage
            exit 0
            ;;
        -q|--quality)
            PLOT_QUALITY="$2"
            case "$PLOT_QUALITY" in
                low) PLOT_DPI="150" ;;
                medium) PLOT_DPI="200" ;;
                high) PLOT_DPI="300" ;;
                ultra) PLOT_DPI="600" ;;
                *) echo -e "${RED}ERROR: Invalid quality: $PLOT_QUALITY${NC}" >&2; exit 1 ;;
            esac
            shift 2
            ;;
        -f|--format)
            PLOT_FORMAT="$2"
            shift 2
            ;;
        -d|--dpi)
            PLOT_DPI="$2"
            shift 2
            ;;
        -t|--threads)
            JULIA_THREADS="$2"
            shift 2
            ;;
        -o|--output)
            CUSTOM_OUTPUT="$2"
            shift 2
            ;;
        -i|--interactive)
            INTERACTIVE="true"
            shift
            ;;
        --no-report)
            GENERATE_REPORT="false"
            shift
            ;;
        --clean)
            CLEAN_RUN="true"
            shift
            ;;
        --backend)
            PLOT_BACKEND="$2"
            shift 2
            ;;
        full|quick|convergence|spatial|stats|performance|custom)
            VIS_TYPE="$1"
            shift
            ;;
        *)
            echo -e "${RED}ERROR: Unknown option: $1${NC}" >&2
            print_usage
            exit 1
            ;;
    esac
done

# Use custom output directory if specified
if [[ -n "$CUSTOM_OUTPUT" ]]; then
    VIS_DIR="$CUSTOM_OUTPUT"
    mkdir -p "$VIS_DIR"
fi

# Main execution
print_banner

echo "🔧 Configuration:"
echo "   • Visualization Type: $VIS_TYPE"
echo "   • Plot Quality: $PLOT_QUALITY (${PLOT_DPI} DPI)"
echo "   • Output Format: $PLOT_FORMAT"
echo "   • Plot Backend: $PLOT_BACKEND"
echo "   • Interactive Mode: $INTERACTIVE"
echo "   • Generate Report: $GENERATE_REPORT"
echo "   • Julia Threads: $JULIA_THREADS"
echo ""

# Check dependencies and setup
check_visualization_deps
clean_previous_results

# Run visualization
run_visualization "$VIS_TYPE"

echo ""
echo -e "${GREEN}🎉 Visualization script completed successfully!${NC}"
echo ""
echo "📋 Summary:"
echo "   📁 Results: $VIS_DIR"
echo "   📝 Logs: $LOG_DIR"
echo "   🎨 Format: $PLOT_FORMAT ($PLOT_DPI DPI)"
echo ""

if [[ "$INTERACTIVE" == "true" ]]; then
    echo -e "${BLUE}💡 Interactive plots are available - use Julia to explore them!${NC}"
fi